package com.jd.genie.agent.agent;

import com.alibaba.fastjson.JSON;
import com.jd.genie.agent.dto.Message;
import com.jd.genie.agent.dto.tool.ToolCall;
import com.jd.genie.agent.dto.tool.ToolChoice;
import com.jd.genie.agent.enums.AgentState;
import com.jd.genie.agent.enums.RoleType;
import com.jd.genie.agent.llm.LLM;
import com.jd.genie.agent.printer.Printer;
import com.jd.genie.agent.tool.BaseTool;
import com.jd.genie.agent.util.FileUtil;
import com.jd.genie.agent.util.SpringContextHolder;
import com.jd.genie.config.GenieConfig;
import com.jd.genie.model.response.AgentResponse;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.ApplicationContext;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;

/**
 * 工具调用代理 - 处理工具/函数调用的基础代理类
 */
@Slf4j
public class ReactImplAgent extends ReActAgent {

    private List<ToolCall> toolCalls;
    private String systemPromptSnapshot;
    private String nextStepPromptSnapshot;

    private final String reactSystemPrompt;

    {
        final String SYSTEM_PROMPT = """
                # 要求
                - 如果回答用户问题时，需要输出文件、输出报告，尽量使用HTML网页报告输出网页版报告，如果用户指定“输出表格”尽量使用excel或者csv输出数据
                - 输出报告前，尽量使用搜索工具搜索
                
                # 解决问题的流程
                请使用交替进行的“思考（Thought）、行动（Action）、观察（Observation）'三个步骤来系统地解决回答任务。
                
                思考：基于当前获得的信息进行推理和反思，明确下一步行动的目标。
                
                行动：用于表示需要调用的工具，每一步行动必须是以下两种之一：
                1、工具调用 [Function Calling]：根据任务需要，确定调用工具。
                2、Finish[答案]：得出明确答案后使用此操作，返回答案并终止任务。
                
                观察：记录前一步行动的结果。
                
                你可以进行多轮推理和检索，但必须严格按照上述格式进行操作，尤其是每一步“行动”只能使用上述两种类型之一。
                
                # 示例
                
                问题：科罗拉多造山带东部区域延伸到的区域的海拔范围是多少?
                
                思考：我需要搜索“科罗拉多造山带”，获取它的概况，特别是东部延伸区域的信息。
                行动：搜索[科罗拉多造山带]
                观察：科罗拉多造山带是科罗拉多及其周边地区造山运动的一段。
                
                思考：这里没有提到东部延伸区域的具体信息，我需要继续查找“东部区域”。
                行动：查找 [东部区域]
                观察：（结果 1 / 1）东部区域延伸至高平原，称为中原造山带。
                
                思考：我已经知道东部区域是高平原，我需要查找高平原的海拔范围。
                行动：搜索  [高平原 海拔]
                观察：高平原的海拔高度从 1800 到 7000 英尺。
                
                思考：我已经得到了答案，可以结束任务。
                行动：Finish[1800 到 7000 英尺]
                # 语言设置
                - 默认工作语言为**中文**，如用户明确指定其他语言，则按用户要求切换。
                - 所有思考、推理与输出均应使用当前工作语言。
                # 当前环境变量
                - 当前日期：{{date}}
                - 可用文件及描述：{{files}}
                - 用户原始任务内容：{{query}}
                现在请回答以下问题：
                """;
        reactSystemPrompt = """
                # 角色
                你是一个超级智能体，名叫Genie。
                
                # 要求
                - 使用 report tool 工具之前，需要获取足够多的信息，先使用搜索工具搜索最新的信息、资讯来进行获取相关信息。
                - 如果回答用户问题时，如果用户没有指定输出格式，使用HTML网页报告输出网页版报告，如果用户指定了输出格式，则按用户指定的格式输出。
                - 如果用户指定输出格式，是指任务的最终输出格式使用该格式输出，中间任务不得使用网页版格式输出，如果用户没有指定输出格式，最后一个任务使用HTML网页报告输出网页版报告。
                - 如果用户指定“输出表格”、“结构化展示”、“结构化输出”或者“抽取相关指标”，尽量使用excel或者csv输出数据；如果已经生成了相应的Excel、csv文件，说明已经满足了“结构化展示”、“结构化输出”等要求。
                - 优先选择合适的工具完成任务，不要重复使用相同工具进行尝试
                
                # 语言要求
                - 所有内容均以 **中文** 输出
                
                # 解决问题的流程
                请使用交替进行的“思考、行动、观察"三个步骤来系统地解决回答任务。
                
                思考：基于当前获得的信息进行推理和反思，明确下一步行动的目标，使用平文本输出，不超过200字。
                
                行动：用于表示需要调用的工具，每一步行动必须是以下两种之一：
                1、工具调用：根据任务需要，确定调用工具。如果用户问题是从上传的文件中直接抽取相关指标，请不要调用code_interpreter工具。
                2、完成：得出明确答案后使用此操作，返回答案并终止任务。
                
                观察：记录前一步行动的结果。
                
                你可以进行多轮推理和检索，但必须严格按照上述格式进行操作，尤其是每一步“行动”只能使用上述两种类型之一。
                
                # 示例
                以下是一些逐步解决问题的示例，仅供参考。
                
                ## 示例 1：
                科罗拉多造山带东部区域延伸到的区域的海拔范围是多少?
                
                思考：了解科罗拉多造山带东部区域延伸到的区域的海拔范围，我需要先明确科罗拉多造山带东部区域延伸到哪些区域，再查找这些区域的海拔范围。第一步，我将通过网络搜索获取科罗拉多造山带东部区域延伸到的区域的相关信息。
                行动：搜索[“科罗拉多造山带概况，特别是东部延伸区域的信息”]
                观察：科罗拉多造山带是科罗拉多及其周边地区造山运动的一段。
                
                思考：通过搜索得知，科罗拉多造山带东部区域延伸至高平原。接下来，我需要搜索高平原的海拔范围。
                行动：搜索[高平原的海拔范围]
                观察：科罗拉多造山带东部区域延伸至高平原，高平原是大平原的一个分区，其海拔从 1800 到 7000 英尺（550 到 2130 米）不等。
                
                思考：我已经得到了答案[1800 到 7000 英尺]，可以结束任务。
                
                ## 示例 2
                分析一下三大电商平台京东、淘宝、拼多多的优劣势
                
                思考：分析京东、淘宝、拼多多三大电商平台的优劣势，我计划先分别搜索各平台优势和劣势的相关信息，再进行整理和分析，最后将结果保存为 HTML 文件。第一步，我需要使用 搜索 工具搜索京东、淘宝、拼多多优势和劣势的相关内容。
                行动：搜索 [搜索京东、淘宝、拼多多优势和劣势的相关内容]
                观察：搜索结果已经保存到文件中。
                
                思考：已获取到京东、淘宝、拼多多优势和劣势的相关信息，接下来我将对这些信息进行整理和分析，形成一份详细的分析报告，并使用工具将输出 HTML 报告文件。
                行动：执行 HTML 报告工具
                观察：已获取到京东、淘宝、拼多多优势和劣势的相关信息，接下来我将对这些信息进行整理和分析，形成一份详细的分析报告。
                
                思考：我已经得到了答案，可以结束任务
                
                ## 示例3
                从上传的文件中抽取指标或者数据，并结构化展示和输出。
                
                思考：我将读取分析好文件内容，抽取相关指标并整理成结构化的表格形式，然后使用文件保存工具将结果保存为csv文件。
                行动：调用 文件工具 读取文件内容
                观察：我已经获取到文件内容，内容中包含Markdown格式的表格
                
                思考：现在提取出文件内容中的表格数据，然后使用文件保存工具将结果保存文件。
                行动：调用 文件工具 保存表格文件文件
                观察：已经抽取表格保存成文件
                
                思考：我已经得到了答案，可以结束任务。
                现在请回答用户问题：
                \s
                # 当前环境变量
                ## 当前日期
                <date>
                {{date}}
                </date>
                
                ## 可用文件及描述
                <files>
                {{files}}\s
                </files>
                
                ## 用户历史对话信息
                <history_dialogue>
                {{history_dialogue}}
                </history_dialogue>
                
                ## 失败处理
                - 不要使用相同入参重复调用失败的工具。
                
                ## 重复处理
                - 应优先利用已有内容，避免重复操作，重复调用相同工具。\s
                \s
                一步一步思考，逐步思考，然后使用工具完成用户的问题或任务。
                """;
    }

    private final String reactNextStepPrompt;

    {
        final String NEXT_STEP_PROMPT = "根据当前状态和可用工具，确定下一步行动，尽可能完成任务\\n\\n输出100字以内的纯文本思考（Reasoning）和反思过程，然后根据思考使用工具来完成当前任务 -判断任务是否已经完成：\\n- 当前任务已完成，则不调用工具。\\n- 当前任务未完成，尝试使用工具完成当前任务，如果尝试潜在能完成任务的工具后，依旧没有办法完成，请通过你过往的知识回答。";
        reactNextStepPrompt = """
                根据当前状态和可用工具，确定下一步行动，根据之前的执行结果，继续完成用户的任务：<task>{{query}}</task>，还需要执行什么工具来继续完成任务。
                -先判断任务是否已经完成：
                - 如果当前任务已完成，则不调用工具。
                - 如果当前任务未完成，尽可能使用工具调用来完成任务。
                
                先输出200字以内的纯文字，（不要重复之前的思考，不能透露代码、链接等。严禁使用Markdown格式输出思考过程，不要重复文件中的内容，仅摘要文件中部分关键内容，不超过200字内容。），再根据任务完成情况使用工具（严禁使用相同入参执行相同的工具，输出相同的文件）来完成任务。（其中，‘工具执行结果：...’是用于标识完成执行工具后得到的内容，你不能重复历史内容，尤其是严禁输出‘工具执行结果’标识。其中，工具执行结果为:null，表示工具执行失败，请不要重复执行失败的工具）
                """;
    }

    // an agent that can execute tool calls.
    public ReactImplAgent(AgentContext context, Printer printer)
    {
        super(context);
        ApplicationContext applicationContext = SpringContextHolder.getApplicationContext();
        GenieConfig genieConfig = applicationContext.getBean(GenieConfig.class);

        StringBuilder toolPrompt = new StringBuilder();
        for (BaseTool tool : context.getToolCollection().getToolMap().values()) {
            toolPrompt.append(String.format("工具名：%s 工具描述：%s\n", tool.getName(), tool.getDescription()));
        }

        setSystemPrompt(reactSystemPrompt
                .replace("{{tools}}", toolPrompt.toString())
                .replace("{{query}}", context.getQuery())
                .replace("{{date}}", context.getDateInfo()));
        setNextStepPrompt(reactNextStepPrompt
                .replace("{{tools}}", toolPrompt.toString())
                .replace("{{query}}", context.getQuery())
                .replace("{{date}}", context.getDateInfo()));

        systemPromptSnapshot = (getSystemPrompt());
        nextStepPromptSnapshot = (getNextStepPrompt());

        setPrinter(printer);
        setLlm(new LLM(genieConfig.getReactModelName(), ""));

        // 初始化工具集合
        availableTools = context.getToolCollection();
        setDigitalEmployeePrompt(genieConfig.getDigitalEmployeePrompt());
    }

    public boolean think()
    {
        // 获取文件内容
        String filesStr = FileUtil.formatFileInfo(context.getProductFiles(), true);
        setSystemPrompt(systemPromptSnapshot.replace("{{files}}", filesStr));
        setNextStepPrompt(nextStepPromptSnapshot.replace("{{files}}", filesStr));

        if (!getMemory().getLastMessage().getRole().equals(RoleType.USER)) {
            Message userMsg = Message.userMessage(getNextStepPrompt(), null);
            getMemory().addMessage(userMsg);
        }
        try {
            // 获取带工具选项的响应
            context.setStreamMessageType("tool_thought");

            LLM.ToolCallResponse response = getLlm().askTool(
                    context,
                    getMemory().getMessages(),
                    Message.systemMessage(getSystemPrompt(), null),
                    availableTools,
                    ToolChoice.AUTO, null, context.getIsStream(), 300,
                    printer);

            toolCalls = (response.getToolCalls());

            // 记录响应信息
            if (!context.getIsStream() && response.getContent() != null && !response.getContent().isEmpty()) {
                printer.send("tool_thought", response.getContent());

            }

            // 创建并添加助手消息
            Message assistantMsg;
            if (response.getToolCalls() != null && !response.getToolCalls().isEmpty()) {
                assistantMsg = Message.fromToolCalls(response.getContent(), response.getToolCalls());
            } else {
                assistantMsg = Message.assistantMessage(response.getContent(), null);
            }
            getMemory().addMessage(assistantMsg);
        } catch (Exception e) {

            log.error("{} react think error", context.getRequestId(), e);
            getMemory().addMessage(Message.assistantMessage(
                    "Error encountered while processing: " + e.getMessage(), null));
            setState(AgentState.FINISHED);
            return false;
        }

        return true;
    }

    public String act()
    {
        if (toolCalls.isEmpty()) {
            setState(AgentState.FINISHED);
            return getMemory().getLastMessage().getContent();
        }
        Integer maxObserve = null;
        // action
        Map<String, String> toolResults = executeTools(toolCalls);
        List<String> results = new ArrayList<>();
        for (ToolCall command : toolCalls) {
            String result = toolResults.get(command.getId());
            if (!Arrays.asList("code_interpreter", "report_tool", "file_tool", "deep_search").contains(command.getFunction().getName())) {
                String toolName = command.getFunction().getName();
                printer.send("tool_result", AgentResponse.ToolResult.builder()
                        .toolName(toolName)
                        .toolParam(JSON.parseObject(command.getFunction().getArguments(), Map.class))
                        .toolResult(result)
                        .build(), null);
            }

            if (maxObserve != null) {
                result = result.substring(0, Math.min(result.length(), maxObserve));
            }

            Message toolMsg = Message.toolMessage(
                    result,
                    command.getId(),
                    null
            );
            getMemory().addMessage(toolMsg);
            results.add(result);
        }

        return String.join("\n\n", results);
    }

}